import random

__author__ = 'Jakub Narloch'

from WEDT.DocumentProcessors import DocumentProcessor
from WEDT.Core import FeatureSelection, DocumentReader
import nltk

class ClassificationResult:
    __documentClasses = list
    __accuracy = 0

    def __init__(self):
        pass

    def get_document_classes(self):
        return self.__documentClasses

    def set_document_classes(self, documentClasses):
        self.__documentClasses = documentClasses

    def get_accuracy(self):
        return self.__accuracy

    def set_accuracy(self, accuracy):
        self.__accuracy = accuracy

class ClassificationAlgorithm:
    def __init__(self):
        pass

    def train(self, data):
        pass

    def classify(self, data):
        pass

class NltkClassificationAlgorithmWrapper(ClassificationAlgorithm):
    __algorithm = None
    __classification = None

    def __init__(self, algorithm=None):
        self.__algorithm = algorithm

    def train(self, data):
        self.__classification = self.__algorithm.train(data)

    def classify(self, data):
        return self.__classification.batch_classify(data)

    def accuracy(self, data):
        return nltk.classify.accuracy(self.__classification, data)

class NltkDecisionTreeClassifier(NltkClassificationAlgorithmWrapper):
    def __init__(self):
        NltkClassificationAlgorithmWrapper.__init__(self, nltk.DecisionTreeClassifier)

    def classify(self, data):
        tmpData = list()
        for d in data:
            tmpData.append(d[0])

        return NltkClassificationAlgorithmWrapper.classify(self, tmpData)

class NltkNaiveBayesClassifier(NltkClassificationAlgorithmWrapper):
    def __init__(self):
        NltkClassificationAlgorithmWrapper.__init__(self, nltk.NaiveBayesClassifier)

    def classify(self, data):
        tmpData = list()
        for d in data:
            tmpData.append(d[0])

        return NltkClassificationAlgorithmWrapper.classify(self, tmpData)

class DocumentClassifier:
    __wordProcessor = None
    __classificationAlgorithm = None
    __trainingSet = list()
    __testSet = list()
    __wordFeatures = None
    __progressListener = None

    __featureCount = 3000

    def __init__(self):
        pass

    def setWordProcessor(self, wordProcessor):
        self.__wordProcessor = wordProcessor

    def setClassificationAlgorithm(self, algorithm):
        self.__classificationAlgorithm = algorithm

    def setTrainingSet(self, path):
        self.__trainingSet = path

    def setTestSet(self, path):
        self.__testSet = path

    def setProgressListener(self, progressListener):
        self.__progressListener = progressListener

    def getProgressListener(self):
        return self.__progressListener

    def setFeatureCount(self, featureCount):
        self.__featureCount = featureCount

    def get_document_features(self, words):
        features = {}
        tmp = set(words)
        for word in self.__wordFeatures:
            features[word] = (word in tmp)
        return features

    def classify(self):
        documentProcessor = DocumentProcessor(DocumentReader(self.__wordProcessor))
        documentProcessor.setProgressListener(self.__progressListener)
        featureSelection = FeatureSelection()
        featureSelection.setFeatureMaxCount(self.__featureCount)
        featureSelection.setProgressListener(self.__progressListener)

        print 'pobieranie slow z plikow treningowych'
        trainingDocuments = list()
        trainingDocuments.append(documentProcessor.processDocuments(self.__trainingSet))
        print ' + zakonczono pobieranie slow'

        print 'wyliczanie wag'
        self.__wordFeatures = featureSelection.calculate(trainingDocuments[0])
        print self.__wordFeatures
        print ' + zakonczono wyliczanie wag'

        train_featuresets = list()
        for doc in trainingDocuments[0]:
            temp = [self.get_document_features(doc.get_words()), doc.get_category()]
            train_featuresets.append(temp)
            #print train_featuresets

        testingDocuments = list()
        testingDocuments.append(documentProcessor.processDocuments(self.__testSet))

        test_featuresets = list()
        for doc in testingDocuments[0]:
            temp = [self.get_document_features(doc.get_words()), doc.get_category()]
            test_featuresets.append(temp)
            # print test_featuresets

        if self.__progressListener is not None:
            self.__progressListener.notifyProgressStarted('Trenowanie algorytmu', 1)

        print 'trenowanie'
        self.__classificationAlgorithm.train(train_featuresets)
        print ' + trenowanie'
        # this is workaround to display the results, should be replaced with some more sophisticate solution
        print 'klasyfikowanie'
        if self.__progressListener is not None:
            self.__progressListener.notifyProgressStarted('Klasyfikacja', 1)

        classResult = self.__classificationAlgorithm.classify(test_featuresets)

        accuracy = self.__classificationAlgorithm.accuracy(test_featuresets)

        result = [(testingDocuments[0][i].get_name(), classResult[i]) for i in range(len(testingDocuments[0]))]
        print result
        print ' + klasyfikowanie'

        if self.__progressListener is not None:
            self.__progressListener.notifyProgressCompleted()

        classificationResult = ClassificationResult()
        classificationResult.set_document_classes(result)
        classificationResult.set_accuracy(accuracy)

        return classificationResult

class RandomForestDocumentClassifier:
    __wordProcessor = None
    __trainingSet = list()
    __testSet = list()
    __wordFeatures = None
    __progressListener = None

    __featureCount = 3000
    __numberOfTrees = 3
    __sampleSize = 2500

    def __init__(self):
        pass

    def setWordProcessor(self, wordProcessor):
        self.__wordProcessor = wordProcessor

    def setTrainingSet(self, path):
        self.__trainingSet = path

    def setTestSet(self, path):
        self.__testSet = path

    def setProgressListener(self, progressListener):
        self.__progressListener = progressListener

    def getProgressListener(self):
        return self.__progressListener

    def get_document_features(self, words):
        features = {}
        tmp = set(words)
        for word in self.__wordFeatures:
            features[word] = (word in tmp)
        return features

    def setFeatureCount(self, featureCount):
        self.__featureCount = featureCount

    def set_number_of_trees(self, numberOfTrees):
        self.__numberOfTrees = numberOfTrees

    def set_sample_size(self, sampleSize):
        self.__sampleSize = sampleSize

    def classify(self):
        documentProcessor = DocumentProcessor(DocumentReader(self.__wordProcessor))
        documentProcessor.setProgressListener(self.__progressListener)
        featureSelection = FeatureSelection()
        featureSelection.setFeatureMaxCount(self.__featureCount)
        featureSelection.setProgressListener(self.__progressListener)

        print 'pobieranie slow z plikow treningowych'
        trainingDocuments = documentProcessor.processDocuments(self.__trainingSet)
        print ' + zakonczono pobieranie slow'

        print 'wyliczanie wag'

        #        self.__wordFeatures = featureSelection.calculate(trainingDocument)
        features = featureSelection.calculate(trainingDocuments)
        print features
        print ' + zakonczono wyliczanie wag'

        testingDocuments = documentProcessor.processDocuments(self.__testSet)

        if self.__progressListener is not None:
            self.__progressListener.notifyProgressStarted('Trenowanie algorytmu', self.__numberOfTrees)

        print 'trenowanie'
        if self.__progressListener is not None:
            self.__progressListener.notifyProgressStarted('Trenowanie algorytmu', self.__numberOfTrees)

        algorithms = list()
        algorithmsFeatures = list()
        for i in range(self.__numberOfTrees):
            self.__wordFeatures = random.sample(features, self.__sampleSize)
            algorithmsFeatures.append(self.__wordFeatures)

            train_featuresets = list()
            for doc in trainingDocuments:
                temp = [self.get_document_features(doc.get_words()), doc.get_category()]
                train_featuresets.append(temp)

            algorithms.append(nltk.DecisionTreeClassifier.train(train_featuresets))

            if self.__progressListener is not None:
                self.__progressListener.notifyProgress(i + 1)

        print 'klasyfikowanie'
        if self.__progressListener is not None:
            self.__progressListener.notifyProgressStarted('Klasyfikacja', self.__numberOfTrees)

        algorithmResults = list()
        for i in range(self.__numberOfTrees):
            self.__wordFeatures = algorithmsFeatures[i]

            test_featuresets = list()
            for doc in testingDocuments:
                temp = [self.get_document_features(doc.get_words()), doc.get_category()]
                test_featuresets.append(temp)
            # print test_featuresets

            tmpData = list()
            for d in test_featuresets:
                tmpData.append(d[0])

            algorithmResults.append(algorithms[i].batch_classify(tmpData))

            if self.__progressListener is not None:
                self.__progressListener.notifyProgress(i + 1)

        classResult = list()
        for i in range(len(algorithmResults[0])):
            classes = dict()
            maxValuesCount = 0
            maxValue = 0
            maxIndex = None
            for j in range(self.__numberOfTrees):
                tmpClass = algorithmResults[j][i]
                if tmpClass in classes:
                    classes[tmpClass] += 1
                else:
                    classes[tmpClass] = 1

                if classes[tmpClass] > maxValue:
                    maxValue = classes[tmpClass]
                    maxValuesCount = 1
                    maxIndex = tmpClass
                elif classes[tmpClass] == maxValue:
                    maxValuesCount += 1

            if maxValuesCount > 1:
                maxIndex = random.sample([key for key in classes if classes[key] == maxValue])

            classResult.append(maxIndex)

        # print classResult

        error = 0
        for i in range(len(classResult)):
            print 'testingDocuments[i].get_category() ', testingDocuments[i].get_category(), ', classResult[i] ', classResult[i]
            if testingDocuments[i].get_category() != classResult[i]:
                error += 1

        accuracy = float(len(classResult) - error ) / float(len(classResult))

        result = [(testingDocuments[i].get_name(), classResult[i]) for i in range(len(testingDocuments))]
        print result
        print ' + klasyfikowanie'

        if self.__progressListener is not None:
            self.__progressListener.notifyProgressCompleted()

        classificationResult = ClassificationResult()
        classificationResult.set_document_classes(result)
        classificationResult.set_accuracy(accuracy)

        return classificationResult